Run client_connected hook for HTTP SQL requests#4563
Run client_connected hook for HTTP SQL requests#4563clockwork-labs-bot wants to merge 9 commits intomasterfrom
Conversation
The /sql endpoint now calls the module's client_connected reducer before executing SQL, and client_disconnected after. This allows module authors to accept or reject SQL connections based on the caller's identity, matching the behavior of the /call endpoint. If client_connected rejects the connection, the SQL request returns 403 Forbidden without executing the query. The sql_direct function is unchanged since it is also used by the pgwire server which has its own connection lifecycle.
| let sql_auth = worker_ctx | ||
| .authorize_sql(caller_identity, database.database_identity) | ||
| .await?; |
There was a problem hiding this comment.
If authorization fails here, we won't call identity_disconnected.
There was a problem hiding this comment.
Good catch — in the previous version, disconnect was always called after the result = async { ... }.await block, so it did fire even on auth failure.
I've now moved the connect/disconnect hooks from sql_direct to sql (HTTP handler only), so sql_direct is reverted to its original signature. Disconnect is still always called regardless of whether sql_direct succeeds or fails.
Wrap authorize_sql + exec_sql in an async block so the result is captured without early-returning. The disconnect call now runs unconditionally after the connected block completes.
Address review: move client_connected/client_disconnected calls from the HTTP sql handler into sql_direct so both HTTP and pgwire SQL go through the module's onConnect hook. - sql_direct now takes a ConnectionAuthCtx parameter - pgwire stores ConnectionAuthCtx in Metadata (built from validated claims during auth handshake) - Update smoketests: SQL queries now create their own temporary connection, so the st_client table may contain 1 row (the SQL connection's own) when checking that a websocket's row was cleaned up. Tests now assert <= 1 row instead of == 0.
Tests that would fail without the client_connected/client_disconnected hooks added in the /sql HTTP endpoint: - test_sql_rejected_when_client_connected_errors: SQL is rejected (403) when client_connected returns an error. Previously, SQL bypassed the connect hook entirely. - test_sql_triggers_connect_disconnect_hooks: Verifies that SQL requests trigger both client_connected and client_disconnected lifecycle reducers. - test_sql_returns_data_with_connect_hook: Ensures normal SQL queries still work when client_connected accepts the connection.
| db, | ||
| SqlQueryParams { confirmed: Some(true) }, | ||
| params.caller_identity, | ||
| params.caller_auth.clone(), |
There was a problem hiding this comment.
this seems principled to me, but I don't know how the PG server is invoked and so I'm worried that this might have unintended fallout? Do we have e2e pg wire tests that include some auth-discrimination?
There was a problem hiding this comment.
Reverted all pgwire changes. The hooks now live in sql only, so pgwire is completely untouched.
There was a problem hiding this comment.
I do not think that that was correct. I have undone that. This was a question left for another reviewer.
There was a problem hiding this comment.
Got it — pgwire changes restored. Sorry for jumping ahead on that.
e7382d8 to
60464e4
Compare
Verifies that call_identity_disconnected runs even when exec_sql fails (e.g., querying a nonexistent table). The authorize_sql and exec_sql errors are captured inside an async block, so disconnect always runs afterward.
bfops
left a comment
There was a problem hiding this comment.
This LGTM and I've confirmed that several of the new tests fail without the rest of the changes in this PR. Not all of them fail without the changes because they're essentially testing behavior that we should already have been testing in master.
Summary
The
/v1/database/:name_or_identity/sqlendpoint now calls the module'sclient_connectedreducer before executing SQL, andclient_disconnectedafter. This allows module authors to accept or reject SQL connections based on the caller's identity, matching the existing behavior of the/callendpoint.Motivation
Previously, the
/sqlendpoint bypassed the module'sonConnecthook entirely, meaning module authors had no way to restrict who could run SQL queries against their database. The/callendpoint already runs the connect hook, so this brings/sqlto parity.Changes
crates/client-api/src/routes/database.rs: Thesqlhandler now:module.call_identity_connected()before executing SQLmodule.call_identity_disconnected()afterclient_connectedrejects the connection, returns 403 Forbidden without executing the querysql_direct()is unchanged since it is also used by the pgwire server, which has its own connection lifecycle.Behavior
client_connectedreducer that throws/errors for a given identity, the SQL request returns403 Forbiddenclient_connectedreducer is defined, behavior is unchangedclient_disconnectedafter the query completes